import numpy as np


# 根据索引i、值value和方向direction分割dataSet，返回大于/小于value的子数据集
def splitContinuousDataSet(dataSet,i,value,direction,center,labels):
    subCenter= {}
    if direction==0:
        subDataSet=dataSet[dataSet[:,i]>value]
        sublabel = labels[dataSet[:,i]>value]
        for d in center.keys():
            if center[d][i]>value:
                subCenter[d]=center[d]
    if direction==1:
        subDataSet=dataSet[dataSet[:,i]<=value]
        sublabel = labels[dataSet[:, i] <= value]
        for d in center.keys():
            if center[d][i]<=value:
                subCenter[d]=center[d]
    endlabel = sublabel[np.isin(sublabel,list(subCenter.keys()))]
    endDataSet = subDataSet[np.isin(sublabel,list(subCenter.keys()))]
    return endDataSet,subCenter,endlabel



def createTree(dataSet,center,labels,featureImport):
    myTree = {}
    # myTree['dataSet'] = dataSet.tolist()
    # myTree['labels'] = labels.tolist()
    children=[]
    if len(center)==1:
        for d in  center.keys():
            center[d]= center[d].tolist()
        myTree['center'] = center
        return myTree
    bestFeat=0
    bestValue=0
    mistake = float('inf')
    for i in range(len(dataSet[0])):
        if(featureImport[i]):
            max = float('-inf')
            min = float('inf')
            for k,v in center.items():
                if v[i] >max:
                    max = v[i]
                if v[i]<min:
                    min = v[i]
            valueList = set(dataSet[:,i])
            for value in valueList:
                if value>min and value < max:
                    greaterDataSet,greaterCenter,greaterlabels= splitContinuousDataSet(dataSet, i, value, 0,center,labels)
                    smallerDataSet,smallerCenter,smallerlabels = splitContinuousDataSet(dataSet, i, value, 1,center,labels)
                    tempmistake = len(labels)-len(greaterlabels)-len(smallerlabels)
                    if(tempmistake<mistake):
                        mistake = tempmistake
                        bestValue =value
                        bestFeat = i

    smallerDataSet, smallerCenter, smallerlabels=splitContinuousDataSet(dataSet, bestFeat, bestValue, 1,center,labels)
    children.append(createTree(smallerDataSet, smallerCenter, smallerlabels,featureImport))
    greaterDataSet, greaterCenter, greaterlabels = splitContinuousDataSet(dataSet, bestFeat, bestValue, 0, center,
                                                                          labels)
    children.append(createTree(greaterDataSet, greaterCenter, greaterlabels,featureImport))
    for d in center.keys():
        center[d] = center[d].tolist()
    myTree['center'] = center
    myTree['feature'] = bestFeat
    myTree['threshold'] = bestValue
    myTree['children'] = children
    myTree['mistake']= mistake
    return myTree

def createClusterTree(dataSet,labels,featureImport):
    for i in range(len(labels)):
        labels[i] = str(labels[i])
    labels = np.array(labels)
    dataSet = np.array(dataSet)
    dataSet = dataSet.astype(float)
    labelSet = set(labels)
    subSets = {}
    subCenter={}
    for d in labelSet:
        labelbool = labels == d
        subSets[d] = dataSet[labelbool]
        subCenter[str(d)] = np.average(subSets[d],axis=0)
    return createTree(dataSet,subCenter,labels,featureImport)


def getOrderFeature(dataSet,labels,center):
    for i in range(len(labels)):
        labels[i] = str(labels[i])
    labels = np.array(labels)
    dataSet = np.array(dataSet)
    dataSet = dataSet.astype(float)
    templabels = labels[np.isin(labels, list(center.keys()))]
    dataSet = dataSet[np.isin(labels, list(center.keys()))]
    labels = templabels
    bestValues=[]
    for i in range(len(dataSet[0])):
        bestValue= 0
        mistake=float('inf')
        max = float('-inf')
        min = float('inf')
        for k,v in center.items():
            if v[i] >max:
                max = v[i]
            if v[i]<min:
                min = v[i]
        valueList = set(dataSet[:,i])
        mistakeValue=[];
        for value in valueList:
            if value>=min and value <= max:
                greaterDataSet,greaterCenter,greaterlabels= splitContinuousDataSet(dataSet, i, value, 0,center,labels)
                smallerDataSet,smallerCenter,smallerlabels = splitContinuousDataSet(dataSet, i, value, 1,center,labels)
                tempmistake = len(labels)-len(greaterlabels)-len(smallerlabels)
                if(tempmistake<mistake):
                    mistake = tempmistake
                    bestValue=value
                mistakeValue.append([value,tempmistake])
        if(mistake==float('inf')):
            mistake = "infinity"
        bestValues.append({'bestValue':bestValue,'mistake':mistake,'mistakelist':mistakeValue,'index':i,'minValue':min,'maxValue':max})

    return bestValues

# 主函数，程序入口
if __name__ == '__main__':
    a=[[1,1],[1,2],[2,1],[2,2],[8,8],[8,9],[9,8],[9,9]]
    b=[0,0,0,0,1,1,1,1]
    print(createClusterTree(a,b))
